from matplotlib import pyplot as plt
import numpy as np
from GLSdata.GLSdataset import idx2class_map
from torch import nn
from torch.autograd import Variable
import torch
import cv2

colors = plt.cm.hsv(np.linspace(0,1,21)).tolist()


def plot_img(model, input_img, annos=None, transforms=None, scale=1):
    '''
    :param model:
    :param input_img:  np array of shape (w,h,3)  RGB
    :param target:    [[xmin1, ymin1, xmax1, ymax1, class1], [xmin2,ymin2,xmax2,ymax2,class2], []...]
                     float relative coord, class not include background starts from 0
    :return:
    '''
    w, h = input_img.shape[:2]

    true_box = []
    img_true = input_img.copy()
    if annos:
        for rxmin, rymin, rxmax, rymax, label in annos:
            xmin, ymin, xmax, ymax, label_name = int(rxmin*w), int(rymin*h), int(rxmax*w), int(rymax*h), idx2class_map[label]

            true_box.append([int(xmin*scale), int(ymin*scale), int(xmax*scale), int(ymax*scale)])

            cv2.rectangle(img_true, (xmin, ymin), (xmax, ymax), (230, 25, 75))
            cv2.putText(img_true, label_name, (xmin, ymin), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (230, 25, 75), 1)



    img_pred = input_img.copy()
    img_ts = input_img.copy()
    if transforms:
        img_ts = transforms(img_ts)
    img_ts = torch.from_numpy(img_ts).float().permute(2,0,1).unsqueeze(0)
    img_ts = Variable(img_ts.cuda())

    detections = model(img_ts)   # (N,1+C,50,5) confidence first
    detections = detections.cpu().data.numpy()[0]

    preds = []
    for i in range(1, detections.shape[0]):
        j = 0
        while detections[i, j, 0] >= 0.6 and j<10:

            score = detections[i, j, 0]
            label_name = idx2class_map[i - 1]
            display_txt = '%s: %.2f' % (label_name, score)

            pt = detections[i, j, 1:]
            pt[0] = int(pt[0] * w)
            pt[1] = int(pt[1] * h)
            pt[2] = int(pt[2] * w)
            pt[3] = int(pt[3] * h)


            cv2.rectangle(img_pred, (pt[0],pt[1]), (pt[2],pt[3]), (230,25,75))
            cv2.putText(img_pred, display_txt, (pt[0], pt[1]), cv2.FONT_HERSHEY_COMPLEX,
                        0.5, (230, 25, 75),1,cv2.LINE_AA)

            j += 1

            preds.append([int(pt[0]*scale), int(pt[1]*scale), int(pt[2]*scale), int(pt[3]*scale)])

    return input_img, img_true, img_pred, preds, true_box
